// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/exec/prefix_sort.h>
#include <pollux/exec/spiller.h>
#include <pollux/exec/window_build.h>

namespace kumo::pollux::exec {
    // Sorts input data of the Window by {partition keys, sort keys}
    // to identify window partitions. This sort fully orders
    // rows as needed for window function computation.
    class SortWindowBuild : public WindowBuild {
    public:
        SortWindowBuild(
            const std::shared_ptr<const core::WindowNode> &node,
            pollux::memory::MemoryPool *pool,
            common::PrefixSortConfig &&prefixSortConfig,
            const common::SpillConfig *spillConfig,
            tsan_atomic<bool> *nonReclaimableSection,
            melon::Synchronized<common::SpillStats> *spillStats);

        ~SortWindowBuild() override {
            pool_->release();
        }

        bool needsInput() override {
            // No partitions are available yet, so can consume input rows.
            return partitionStartRows_.size() == 0;
        }

        void addInput(RowVectorPtr input) override;

        void spill() override;

        std::optional<common::SpillStats> spilledStats() const override;

        void noMoreInput() override;

        bool hasNextPartition() override;

        std::shared_ptr<WindowPartition> nextPartition() override;

    private:
        void ensureInputFits(const RowVectorPtr &input);

        void ensureSortFits();

        void setupSpiller();

        // Main sorting function loop done after all input rows are received
        // by WindowBuild.
        void sortPartitions();

        // Function to compute the partitionStartRows_ structure.
        // partitionStartRows_ is vector of the starting rows index
        // of each partition in the data. This is an auxiliary
        // structure that helps simplify the window function computations.
        void computePartitionStartRows();

        // Find the next partition start row from start.
        vector_size_t findNextPartitionStartRow(vector_size_t start);

        // Reads next partition from spilled data into 'data_' and 'sortedRows_'.
        void loadNextPartitionFromSpill();

        const size_t numPartitionKeys_;

        // Compare flags for partition and sorting keys. Compare flags for partition
        // keys are set to default values. Compare flags for sorting keys match
        // sorting order specified in the plan node.
        //
        // Used to sort 'data_' while spilling and in Prefix sort.
        const std::vector<CompareFlags> compareFlags_;

        memory::MemoryPool *const pool_;

        // Config for Prefix-sort.
        const common::PrefixSortConfig prefixSortConfig_;

        melon::Synchronized<common::SpillStats> *const spillStats_;

        // allKeyInfo_ is a combination of (partitionKeyInfo_ and sortKeyInfo_).
        // It is used to perform a full sorting of the input rows to be able to
        // separate partitions and sort the rows in it. The rows are output in
        // this order by the operator.
        std::vector<std::pair<column_index_t, core::SortOrder> > allKeyInfo_;

        // Vector of pointers to each input row in the data_ RowContainer.
        // The rows are sorted by partitionKeys + sortKeys. This total
        // ordering can be used to split partitions (with the correct
        // order by) for the processing.
        std::vector<char *, memory::StlAllocator<char *> > sortedRows_;

        // This is a vector that gives the index of the start row
        // (in sortedRows_) of each partition in the RowContainer data_.
        // This auxiliary structure helps demarcate partitions.
        std::vector<vector_size_t, memory::StlAllocator<vector_size_t> >
        partitionStartRows_;

        // Current partition being output. Used to construct WindowPartitions
        // during resetPartition.
        vector_size_t currentPartition_ = -1;

        // Spiller for contents of the 'data_'.
        std::unique_ptr<SortInputSpiller> spiller_;

        // Used to sort-merge spilled data.
        std::unique_ptr<TreeOfLosers<SpillMergeStream> > merge_;
    };
} // namespace kumo::pollux::exec
